package org.basis.enhance.async.init.executor;

import org.basis.enhance.async.init.constants.AsyncInitConstants;
import org.basis.enhance.async.init.exception.BasisAsyncInitException;
import org.enhance.common.executor.factory.BasisThreadFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.env.Environment;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Future;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 异步初始化线程池
 *
 * @author wenpan 2023/6/11 4:45 下午
 */
public class AsyncTaskExecutor {

    private final static Logger logger = LoggerFactory.getLogger(AsyncTaskExecutor.class);

    /**
     * CPU核数
     */
    protected static final int CPU_COUNT = Runtime.getRuntime().availableProcessors();
    /**
     * 线程池引用
     */
    protected static final AtomicReference<ThreadPoolExecutor> THREAD_POOL_REF = new AtomicReference<>();
    /**
     * 线程池执行结果future
     */
    protected static final List<Future<?>> FUTURES = new ArrayList<>();
    /**
     * 异步初始化是否执行完毕
     */
    protected static final AtomicBoolean STARTED = new AtomicBoolean(false);

    public static Future<?> submitTask(Environment environment, Runnable runnable) {
        if (THREAD_POOL_REF.get() == null) {
            ThreadPoolExecutor threadPoolExecutor = createThreadPoolExecutor(environment);
            // cas
            boolean success = THREAD_POOL_REF.compareAndSet(null, threadPoolExecutor);
            // 如果 THREAD_POOL_REF 已经被别的线程设置了，则关闭刚创建的线程池
            if (!success) {
                threadPoolExecutor.shutdown();
            }
        }
        // 提交任务到线程池
        Future<?> future = THREAD_POOL_REF.get().submit(runnable);
        FUTURES.add(future);
        return future;
    }

    /**
     * Create thread pool to execute async init task
     */
    private static ThreadPoolExecutor createThreadPoolExecutor(Environment environment) {
        int threadPoolCoreSize = CPU_COUNT + 1;
        String coreSizeStr = environment.getProperty(AsyncInitConstants.ASYNC_INIT_BEAN_CORE_SIZE);
        if (coreSizeStr != null) {
            threadPoolCoreSize = Integer.parseInt(coreSizeStr);
        }

        int threadPoolMaxSize = CPU_COUNT + 1;
        String maxSizeStr = environment.getProperty(AsyncInitConstants.ASYNC_INIT_BEAN_MAX_SIZE);
        if (maxSizeStr != null) {
            threadPoolMaxSize = Integer.parseInt(maxSizeStr);
        }

        logger.info(String.format("create async-init-bean thread pool, corePoolSize: %d, maxPoolSize: %d.",
                threadPoolCoreSize, threadPoolMaxSize));
        return new ThreadPoolExecutor(
                threadPoolCoreSize,
                threadPoolMaxSize,
                30,
                TimeUnit.SECONDS,
                new SynchronousQueue<>(),
                BasisThreadFactory.create("basis-async-init", true),
                new ThreadPoolExecutor.CallerRunsPolicy());
    }

    /**
     * 确保异步init task已经正常finish
     */
    public static void ensureAsyncTasksFinish() {
        for (Future<?> future : FUTURES) {
            try {
                future.get();
            } catch (Throwable e) {
                throw new BasisAsyncInitException("异步初始化失败.", e);
            }
        }

        STARTED.set(true);
        FUTURES.clear();
        if (THREAD_POOL_REF.get() != null) {
            // 关闭并释放线程池
            synchronized (AsyncTaskExecutor.class) {
                if (THREAD_POOL_REF.get() != null) {
                    THREAD_POOL_REF.get().shutdown();
                    THREAD_POOL_REF.set(null);
                }
            }
        }
    }

    /**
     * 是否启动完毕
     */
    public static boolean isStarted() {
        return STARTED.get();
    }
}